// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/json/chunker.h>

#include <algorithm>
#include <string_view>
#include <utility>
#include <vector>

#include <nebula/json/rapidjson_defs.h>
#include <merak/json.h>

#include <nebula/core/buffer.h>
#include <nebula/json/options.h>
#include <turbo/log/logging.h>

using std::string_view;

namespace nebula::json {

    static size_t ConsumeWhitespace(string_view view) {
#ifdef RAPIDJSON_SIMD
        auto data = view.data();
        auto nonws_begin = merak::json::SkipWhitespace_SIMD(data, data + view.size());
        return nonws_begin - data;
#else
        auto ws_count = view.find_first_not_of(" \t\r\n");
        if (ws_count == string_view::npos) {
            return view.size();
        } else {
            return ws_count;
        }
#endif
    }

    class MultiStringStream {
    public:
        using char_type = char;

        explicit MultiStringStream(std::vector<string_view> strings)
                : strings_(std::move(strings)) {
            std::reverse(strings_.begin(), strings_.end());
        }

        explicit MultiStringStream(const BufferVector &buffers) : strings_(buffers.size()) {
            for (size_t i = 0; i < buffers.size(); ++i) {
                strings_[i] = string_view(*buffers[i]);
            }
            std::reverse(strings_.begin(), strings_.end());
        }

        char peek() const {
            if (strings_.size() == 0) return '\0';
            return strings_.back()[0];
        }

        char get() {
            if (strings_.size() == 0) return '\0';
            char taken = strings_.back()[0];
            if (strings_.back().size() == 1) {
                strings_.pop_back();
            } else {
                strings_.back() = strings_.back().substr(1);
            }
            ++index_;
            return taken;
        }

        size_t tellg() { return index_; }

    private:
        size_t index_ = 0;
        std::vector<string_view> strings_;
    };

    template<typename Stream>
    static size_t ConsumeWholeObject(Stream &&stream) {
        static constexpr unsigned parse_flags = merak::json::kParseIterativeFlag |
                                                merak::json::kParseStopWhenDoneFlag |
                                                merak::json::kParseNumbersAsStringsFlag;
        merak::json::BaseReaderHandler<merak::json::UTF8<>> handler;
        merak::json::Reader reader;
        // parse a single JSON object
        switch (reader.parse<parse_flags>(stream, handler).Code()) {
            case merak::json::kParseErrorNone:
                return stream.tellg();
            case merak::json::kParseErrorDocumentEmpty:
                return 0;
            default:
                return string_view::npos;
        }
    }

    namespace {

    // A BoundaryFinder implementation that assumes JSON objects can contain raw newlines,
    // and uses actual JSON parsing to delimit them.
        class ParsingBoundaryFinder : public BoundaryFinder {
        public:
            turbo::Status FindFirst(string_view partial, string_view block, int64_t *out_pos) override {
                auto length = ConsumeWholeObject(MultiStringStream({partial, block}));
                if (length == string_view::npos) {
                    *out_pos = -1;
                } else if (TURBO_UNLIKELY(length < partial.size())) {
                    return turbo::invalid_argument_error("JSON chunk error: invalid data at end of document");
                } else {
                            DKCHECK_LE(length, partial.size() + block.size());
                    *out_pos = static_cast<int64_t>(length - partial.size());
                }
                return turbo::OkStatus();
            }

            turbo::Status FindLast(std::string_view block, int64_t *out_pos) override {
                const size_t block_length = block.size();
                size_t consumed_length = 0;
                while (consumed_length < block_length) {
                    merak::json::MemoryStream ms(reinterpret_cast<const char *>(block.data()), block.size());
                    using InputStream = merak::json::EncodedInputStream<merak::json::UTF8<>, merak::json::MemoryStream>;
                    auto length = ConsumeWholeObject(InputStream(ms));
                    if (length == string_view::npos || length == 0) {
                        // found incomplete object or block is empty
                        break;
                    }
                    consumed_length += length;
                    block = block.substr(length);
                }
                if (consumed_length == 0) {
                    *out_pos = -1;
                } else {
                    consumed_length += ConsumeWhitespace(block);
                            DKCHECK_LE(consumed_length, block_length);
                    *out_pos = static_cast<int64_t>(consumed_length);
                }
                return turbo::OkStatus();
            }

            turbo::Status FindNth(std::string_view partial, std::string_view block, int64_t count,
                           int64_t *out_pos, int64_t *num_found) override {
                return turbo::unimplemented_error("ParsingBoundaryFinder::FindNth");
            }
        };

    }  // namespace

    std::unique_ptr<Chunker> MakeChunker(const ParseOptions &options) {
        std::shared_ptr<BoundaryFinder> delimiter;
        if (options.newlines_in_values) {
            delimiter = std::make_shared<ParsingBoundaryFinder>();
        } else {
            delimiter = MakeNewlineBoundaryFinder();
        }
        return std::make_unique<Chunker>(std::move(delimiter));
    }

}  // namespace nebula::json
